from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
import base64
import os
class Password:
    def __init__(self, password, is_encrypt=True):
        crypt_sm4 = CryptSM4()
        if is_encrypt:
            self.encrypt_pwd = password
            self.base64pwd = self.__urlsafe_decode_from_base64(self.encrypt_pwd)
            key = self.base64pwd[:16]
            encrypt_pwd = self.base64pwd[16:]
            crypt_sm4.set_key(key, SM4_DECRYPT)
            self.real_pwd = crypt_sm4.crypt_ecb(encrypt_pwd).decode('utf-8')
        else:
            self.real_pwd = password
            key = os.urandom(8).hex().encode('utf-8')
            crypt_sm4.set_key(key, SM4_ENCRYPT)
            tmp_encrypt_pwd = crypt_sm4.crypt_ecb(self.real_pwd.encode('utf-8'))
            self.encrypt_pwd = self.__urlsafe_encode_to_base64(key + tmp_encrypt_pwd)

    def __str__(self):
        return self.real_pwd

    def encrypt(self):
        return self.encrypt_pwd 

    # URL 安全的 Base64 编码（替换 + 和 /，删除填充 =）
    def __urlsafe_encode_to_base64(self, data):
        if isinstance(data, str):
            data = data.encode('utf-8')
        encoded_bytes = base64.urlsafe_b64encode(data)
        return encoded_bytes.decode('ascii').rstrip('=')  # 删除填充字符

    # URL 安全的 Base64 解码
    def __urlsafe_decode_from_base64(self, base64_str):
        # 添加缺失的填充字符
        padding_needed = 4 - (len(base64_str) %
                              4) if len(base64_str) % 4 else 0
        base64_str = base64_str + '=' * padding_needed

        if isinstance(base64_str, str):
            base64_str = base64_str.encode('ascii')
        decoded_bytes = base64.urlsafe_b64decode(base64_str)
        return decoded_bytes


if __name__ == "__main__":
    pwd1 = Password("Kermit#@3211", is_encrypt=False)
    print(pwd1)
    print(pwd1.encrypt())
    pwd2 = Password("Kermit#@3211", is_encrypt=False)
    print(pwd2.encrypt())    
    pwd3 = Password("Kermit#@3211", is_encrypt=False)
    print(pwd3.encrypt())    
    pwd4 = Password("Kermit#@3211", is_encrypt=False)
    print(pwd4.encrypt())    
    newpwd = Password(pwd1.encrypt(), is_encrypt=True)
    print(newpwd)    
    newpwd = Password(pwd2.encrypt(), is_encrypt=True)
    print(newpwd)    
    newpwd = Password(pwd3.encrypt(), is_encrypt=True)
    print(newpwd)    
    newpwd = Password(pwd4.encrypt(), is_encrypt=True)
    print(newpwd)